import os
import json
from shutil import copy, rmtree
import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
from model import LeNet5
# import cv2
import numpy as np
import io

def get_image(patient):
    img_path = "./pridiction/DMQ/Axial/"+patient+ '.png'
    img = Image.open(img_path)

    result = main(patient)
    prediction = result['prediction']
    prob = result['probability']
    original_class = result['original']
    print_res = "patient:{}     original:{}     predict: {}     prob: {:.3}".format(patient,original_class,prediction,
                                                 prob)

    fig = plt.figure()
    plt.axis('off')
    plt.imshow(img)
    plt.title(print_res)
    
    canvas = fig.canvas
    buffer = io.BytesIO()
    canvas.print_png(buffer)
    data = buffer.getvalue()
    buffer.close()
    return data


def main(img_path):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    data_transform = transforms.Compose(
        [transforms.ToTensor()])
    assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
    img = Image.open(img_path)
    # [N, C, H, W]
    img = data_transform(img)
    # expand batch dimension
    img = torch.unsqueeze(img, dim=0)
    # read class_indict
    json_path = './class_indices.json'
    assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
    json_file = open(json_path, "r")
    class_indict = json.load(json_file)
    # create model
    model = LeNet5(num_classes=2).to(device)
    # load model weights
    weights_path = "./LeNet5.pth"
    assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
    model.load_state_dict(torch.load(weights_path,map_location='cpu'))
    model.eval()
    with torch.no_grad():
        # predict class
        output = torch.squeeze(model(img.to(device))).cpu()
        predict = torch.softmax(output, dim=0)
        predict_cla = torch.argmax(predict).numpy()

    prediction = class_indict[str(predict_cla)]
    probability = "{:.3}".format(predict[predict_cla].numpy())

    filename = os.path.basename(img_path)
    patient = filename.split('.')[0]
    original_class = get_original_class(patient)

    result = {
        "patient":patient,
        "original":original_class,
        "prediction":prediction,
        "probability":probability
    }

    return result
    

def get_original_class(patient):
    seob0 = ['10752368_1', '10965468_6', '10966721_7', '11181780_8', '11194853_9', '11268900_11', '11392444_12',
             '11458341_13', '11478616_15', '11488504_17', '11509477_18', '11514232_19', '11527001_21', '11533753_22', '11547091_24', '11549625_25',
             '11550347_26', '11562531_29', '11577926_31', '11636923_41']
    seob1 = ['10895033_2', '10907994_3', '10909281_4', '10962827_5', '11260577_10', '11486292_16', '11539683_23', '11556440_28', '11575532_30', '11580722_32', '11583688_33',
             '11590991_34', '11596400_35', '11608853_37', '11609358_38', '11635851_39', '11636923_40', '11635851_42']

    ################################################
    # 84:
    # high:low=17:10
    s840 = ['10007609', '10816997', '11591806', '11107779', '11392444', '11446761', '11612892', '11485403', '10787472', '10967130']
    s841 = ['10605374', '10948373', '10850673', '11040843', '11062316', '10960862', '11480762', '11536646', '11538428', '11566602',
            '10705030', '10009051', '10659923', '10759978', '10672345', '11020305', '11078608']

    ################################################
    # 63:
    # high:low=28:24
    s630 = ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21', '53', '59', '62']
    s631 = ['22', '23', '24', '25', '26', '27', '28', '29', '31', '33', '34', '35', '36', '37', '38', '39', '40', '41', '42', '43', '45', '46', '50', '54', '55', '56', '57', '61']

    low = s630 + s840 + seob0
    high = s631 + s841 + seob1

    original_class = 'low'
    if patient in high:
        original_class='high'
    return original_class


# 错误 26/117
# 正确率 77.78%

"""
patient:1     original:low     pridict: low     prob: 0.739
patient:2     original:low     pridict: low     prob: 0.746
patient:3     original:low     pridict: low     prob: 0.692
patient:4     original:low     pridict: low     prob: 0.679
patient:5     original:low     pridict: low     prob: 0.711
patient:6     original:low     pridict: low     prob: 0.736
patient:7     original:low     pridict: low     prob: 0.76
patient:8     original:low     pridict: low     prob: 0.679
patient:9     original:low     pridict: low     prob: 0.513
patient:10     original:low     pridict: low     prob: 0.541
patient:11     original:low     pridict: low     prob: 0.716
patient:12     original:low     pridict: low     prob: 0.773
patient:13     original:low     pridict: high     prob: 0.587
patient:14     original:low     pridict: low     prob: 0.705
patient:15     original:low     pridict: low     prob: 0.661
patient:16     original:low     pridict: low     prob: 0.712
patient:17     original:low     pridict: low     prob: 0.624
patient:18     original:low     pridict: low     prob: 0.753
patient:19     original:low     pridict: low     prob: 0.728
patient:20     original:low     pridict: low     prob: 0.703
patient:21     original:low     pridict: low     prob: 0.579
patient:53     original:low     pridict: low     prob: 0.747
patient:59     original:low     pridict: low     prob: 0.736
patient:62     original:low     pridict: low     prob: 0.733
patient:10007609     original:low     pridict: low     prob: 0.646
patient:10816997     original:low     pridict: low     prob: 0.671
patient:11591806     original:low     pridict: low     prob: 0.703
patient:11107779     original:low     pridict: high     prob: 0.716
patient:11392444     original:low     pridict: low     prob: 0.715
patient:11446761     original:low     pridict: high     prob: 0.671
patient:11612892     original:low     pridict: low     prob: 0.718
patient:11485403     original:low     pridict: low     prob: 0.697
patient:10787472     original:low     pridict: low     prob: 0.691
patient:10967130     original:low     pridict: high     prob: 0.557
patient:10752368_1     original:low     pridict: low     prob: 0.705
patient:10965468_6     original:low     pridict: low     prob: 0.716
patient:10966721_7     original:low     pridict: low     prob: 0.735
patient:11181780_8     original:low     pridict: low     prob: 0.724
patient:11194853_9     original:low     pridict: high     prob: 0.513
patient:11268900_11     original:low     pridict: low     prob: 0.524
patient:11392444_12     original:low     pridict: low     prob: 0.706
patient:11458341_13     original:low     pridict: low     prob: 0.696
patient:11478616_15     original:low     pridict: low     prob: 0.751
patient:11488504_17     original:low     pridict: high     prob: 0.645
patient:11509477_18     original:low     pridict: low     prob: 0.706
patient:11514232_19     original:low     pridict: low     prob: 0.757
patient:11527001_21     original:low     pridict: low     prob: 0.752
patient:11533753_22     original:low     pridict: low     prob: 0.668
patient:11547091_24     original:low     pridict: low     prob: 0.563
patient:11549625_25     original:low     pridict: low     prob: 0.565
patient:11550347_26     original:low     pridict: high     prob: 0.536
patient:11562531_29     original:low     pridict: low     prob: 0.736
patient:11577926_31     original:low     pridict: low     prob: 0.699
patient:11636923_41     original:low     pridict: low     prob: 0.706


###


patient:22     original:high     pridict: high     prob: 0.68
patient:23     original:high     pridict: high     prob: 0.919
patient:24     original:high     pridict: high     prob: 0.77
patient:25     original:high     pridict: low     prob: 0.716
patient:26     original:high     pridict: high     prob: 0.515
patient:27     original:high     pridict: low     prob: 0.567
patient:28     original:high     pridict: low     prob: 0.707
patient:29     original:high     pridict: high     prob: 0.831
patient:31     original:high     pridict: high     prob: 0.699
patient:33     original:high     pridict: high     prob: 0.637
patient:34     original:high     pridict: high     prob: 0.899
patient:35     original:high     pridict: high     prob: 0.628
patient:36     original:high     pridict: high     prob: 0.759
patient:37     original:high     pridict: high     prob: 0.827
patient:38     original:high     pridict: high     prob: 0.722
patient:39     original:high     pridict: high     prob: 0.841
patient:40     original:high     pridict: low     prob: 0.582
patient:41     original:high     pridict: low     prob: 0.64
patient:42     original:high     pridict: high     prob: 0.799
patient:43     original:high     pridict: high     prob: 0.691
patient:45     original:high     pridict: high     prob: 0.66
patient:46     original:high     pridict: high     prob: 0.828
patient:50     original:high     pridict: high     prob: 0.697
patient:54     original:high     pridict: high     prob: 0.884
patient:55     original:high     pridict: low     prob: 0.714
patient:56     original:high     pridict: high     prob: 0.534
patient:57     original:high     pridict: low     prob: 0.746
patient:61     original:high     pridict: high     prob: 0.614
patient:10605374     original:high     pridict: low     prob: 0.609
patient:10948373     original:high     pridict: high     prob: 0.739
patient:10850673     original:high     pridict: low     prob: 0.718
patient:11040843     original:high     pridict: high     prob: 0.597
patient:11062316     original:high     pridict: low     prob: 0.738
patient:10960862     original:high     pridict: high     prob: 0.691
patient:11480762     original:high     pridict: high     prob: 0.775
patient:11536646     original:high     pridict: high     prob: 0.826
patient:11538428     original:high     pridict: high     prob: 0.634
patient:11566602     original:high     pridict: low     prob: 0.718
patient:10705030     original:high     pridict: low     prob: 0.695
patient:10009051     original:high     pridict: high     prob: 0.799
patient:10659923     original:high     pridict: high     prob: 0.633
patient:10759978     original:high     pridict: high     prob: 0.752
patient:10672345     original:high     pridict: low     prob: 0.717
patient:11020305     original:high     pridict: high     prob: 0.714
patient:11078608     original:high     pridict: high     prob: 0.756
patient:10895033_2     original:high     pridict: high     prob: 0.791
patient:10907994_3     original:high     pridict: high     prob: 0.831
patient:10909281_4     original:high     pridict: low     prob: 0.694
patient:10962827_5     original:high     pridict: low     prob: 0.607
patient:11260577_10     original:high     pridict: high     prob: 0.803
patient:11486292_16     original:high     pridict: low     prob: 0.704
patient:11539683_23     original:high     pridict: high     prob: 0.701
patient:11556440_28     original:high     pridict: high     prob: 0.573
patient:11575532_30     original:high     pridict: high     prob: 0.641
patient:11580722_32     original:high     pridict: high     prob: 0.617
patient:11583688_33     original:high     pridict: high     prob: 0.636
patient:11590991_34     original:high     pridict: low     prob: 0.728
patient:11596400_35     original:high     pridict: high     prob: 0.649
patient:11608853_37     original:high     pridict: low     prob: 0.665
patient:11609358_38     original:high     pridict: high     prob: 0.894
patient:11635851_39     original:high     pridict: high     prob: 0.731
patient:11636923_40     original:high     pridict: high     prob: 0.592
patient:11635851_42     original:high     pridict: low     prob: 0.655
"""